import random
import datetime
import tarfile
import os.path
import tempfile
import shutil
from pathlib import Path
from io import BytesIO

from slugify import slugify
from docker import DockerClient
from docker.errors import ImageNotFound

from tinyops.library.helper import detect_shell
from tinyops.library.task import TaskLogger, TaskEnv
from tinyops.library.config import DockerConfig

from .helper import ExportHelper


class ContainerExecutor:
    def __init__(
        self,
        server: DockerConfig,
        image: str,
        env: TaskEnv,
        timeout: int,
        logger: TaskLogger,
    ):
        self.server = server
        self.image = image
        self.name = f'{server.name}+{image}'
        self._log = logger.target(self.name).log
        self.env = env if env is not None else TaskEnv()
        self.timeout = timeout
        self.shell = None
        self.client = DockerClient(base_url=server.host, timeout=self.timeout)
        self.container = None
        self.container_name = None
        self._export_helper = ExportHelper(
            env=self.env,
            logger=logger.target(self.name),
            handler=self._export_handler,
        )
        self._inited = False

    def _read_chunks(self, chunks):
        fileobj = BytesIO()
        for chunk in chunks:
            fileobj.write(chunk)
        fileobj.seek(0)
        tarobj = tarfile.open(fileobj=fileobj)
        return tarobj

    def _read_file_from_container(self, filepath: str) -> str:
        chunks, stat = self.container.get_archive(filepath)
        tarobj = self._read_chunks(chunks)
        fileobj = tarobj.extractfile(stat['name'])
        content = fileobj.read()
        if isinstance(content, bytes):
            content = content.decode('utf-8')
        return content

    def _auto_detect_shell(self):
        content = self._read_file_from_container('/etc/shells')
        return detect_shell(content)

    def _gen_container_name(self):
        t = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
        r = '{:03d}'.format(random.randint(0, 999))
        return slugify(f'tinyops_{self.name}_{t}_{r}')

    def init(self):
        if self._inited:
            return
        self._inited = True
        try:
            self.client.images.get(self.image)
        except ImageNotFound:
            self._log(f'pull image {self.image}')
            self.client.images.pull(self.image)
        command = '/bin/sh -c "sleep {}"'.format(self.timeout)
        self.container_name = self._gen_container_name()
        self.container = self.client.containers.run(
            self.image, command, name=self.container_name,
            detach=True, auto_remove=True)
        self._log('container {} {} started'.format(self.container.id[:6], self.container_name))
        if not self.shell:
            self.shell = self._auto_detect_shell()

    def close(self):
        if self.container:
            self.container.remove(force=True)
            self._log('container {} {} removed'.format(self.container.id[:6], self.container_name))
        self.container = None

    def _export_handler(self, export_file):
        content = self._read_file_from_container(export_file)
        clean_command = f'rm -f {export_file}'
        exit_code, output = self.container.exec_run(clean_command)
        if str(exit_code) != '0':
            self._log(f'Warning: delete temp file failed: {export_file}')
        return content

    def do_run(self, script: str, exports: list = None):
        self._log(script, sep='>')
        script = 'set -e;\n' + script
        script = self._export_helper.before_run(script, exports)
        script = '{} -c "{}"'.format(self.shell, script)
        # Get exit code of command from exec_run(stream=True)
        # https://github.com/docker/docker-py/issues/1989
        exec_handler = self.client.api.exec_create(
            self.container.id, script, tty=True, environment=self.env.value())
        exec_id = exec_handler['Id']
        output_gen = self.client.api.exec_start(exec_id, tty=True, stream=True)
        outputs = []
        for chunk in output_gen:
            block = chunk.decode()
            outputs.append(block)
            self._log(block)
        exit_code = self.client.api.exec_inspect(exec_id)['ExitCode']
        self._log(exit_code=exit_code)
        if not self._export_helper.after_run(outputs):
            return False
        return str(exit_code) == '0'

    def do_get(self, remote_path, local_path):
        self._log(f'## GET {remote_path} {local_path}', sep='>')
        chunks, stat = self.container.get_archive(remote_path)
        remote_name = stat['name']
        tarobj = self._read_chunks(chunks)
        # docker cp command is hard to implement, so use temp dir + cp to get same behavior.
        temp_dir = Path(tempfile.mkdtemp('tinyops_container_executor'))
        try:
            tarobj.extractall(temp_dir)
            shutil.move(temp_dir / remote_name, local_path)
        finally:
            shutil.rmtree(temp_dir)
        return True

    def do_put(self, local_path, remote_path):
        self._log(f'## PUT {local_path} {remote_path}', sep='>')
        fileobj = BytesIO()
        tarobj = tarfile.open(fileobj=fileobj, mode='w')
        arcname = os.path.basename(remote_path)
        try:
            tarobj.add(local_path, arcname=arcname)
        finally:
            tarobj.close()
        dstname = os.path.dirname(remote_path)
        ok = self.container.put_archive(dstname, fileobj.getvalue())
        return ok
